"""
简单JWT替代方案 - 替代python-jose
"""
import base64
import json
import time
import hashlib
import hmac
import jwt
import logging
from datetime import datetime, timedelta
from typing import Optional, Dict, Any
from jwt.exceptions import PyJWTError

logger = logging.getLogger(__name__)

SECRET_KEY = "your-secret-key-here"  # 在生产环境中应该从配置文件读取
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 60 * 24 * 7  # 7天
REFRESH_TOKEN_EXPIRE_DAYS = 30  # 30天

class JWTError(Exception):
    """JWT错误"""
    pass

def encode(payload: dict, key: str = SECRET_KEY, algorithm: str = "HS256") -> str:
    """编码JWT token"""
    # Header
    header = {"alg": algorithm, "typ": "JWT"}
    header_encoded = base64.urlsafe_b64encode(
        json.dumps(header, separators=(',', ':')).encode()
    ).decode().rstrip('=')
    
    # Payload
    payload_copy = payload.copy()
    if 'exp' not in payload_copy:
        payload_copy['exp'] = int(time.time()) + 3600  # 默认1小时过期
    
    payload_encoded = base64.urlsafe_b64encode(
        json.dumps(payload_copy, separators=(',', ':')).encode()
    ).decode().rstrip('=')
    
    # Signature
    message = f"{header_encoded}.{payload_encoded}"
    signature = base64.urlsafe_b64encode(
        hmac.new(key.encode(), message.encode(), hashlib.sha256).digest()
    ).decode().rstrip('=')
    
    return f"{header_encoded}.{payload_encoded}.{signature}"

def decode(token: str, key: str = SECRET_KEY, algorithms: list = ["HS256"]) -> dict:
    """解码JWT token"""
    try:
        parts = token.split('.')
        if len(parts) != 3:
            raise JWTError("Invalid token format")
        
        header_encoded, payload_encoded, signature_encoded = parts
        
        # 验证签名
        message = f"{header_encoded}.{payload_encoded}"
        expected_signature = base64.urlsafe_b64encode(
            hmac.new(key.encode(), message.encode(), hashlib.sha256).digest()
        ).decode().rstrip('=')
        
        if signature_encoded != expected_signature:
            raise JWTError("Invalid signature")
        
        # 解码payload
        # 添加padding
        payload_encoded += '=' * (4 - len(payload_encoded) % 4)
        payload_json = base64.urlsafe_b64decode(payload_encoded).decode()
        payload = json.loads(payload_json)
        
        # 检查过期时间
        if 'exp' in payload and payload['exp'] < time.time():
            raise JWTError("Token has expired")
        
        return payload
        
    except Exception as e:
        raise JWTError(f"Token decode error: {e}")

def create_access_token(subject: str) -> str:
    """创建访问令牌"""
    try:
        expire = datetime.utcnow() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
        
        payload = {
            "sub": str(subject),
            "exp": expire,
            "iat": datetime.utcnow(),
            "type": "access"
        }
        
        logger.info(f"创建访问令牌: subject={subject}")
        logger.info(f"payload: {payload}")
        logger.info(f"jwt module: {jwt}")
        logger.info(f"jwt.encode: {jwt.encode}")
        
        result = jwt.encode(payload, SECRET_KEY, algorithm=ALGORITHM)
        logger.info(f"JWT创建成功: {result[:20]}...")
        return result
        
    except Exception as e:
        logger.error(f"创建访问令牌失败: {str(e)}")
        logger.error(f"jwt类型: {type(jwt)}")
        raise

def create_refresh_token(subject: str) -> str:
    """创建刷新令牌"""
    try:
        expire = datetime.utcnow() + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
        
        payload = {
            "sub": str(subject),
            "exp": expire,
            "iat": datetime.utcnow(),
            "type": "refresh"
        }
        
        logger.info(f"创建刷新令牌: subject={subject}")
        return jwt.encode(payload, SECRET_KEY, algorithm=ALGORITHM)
        
    except Exception as e:
        logger.error(f"创建刷新令牌失败: {str(e)}")
        raise

def verify_token(token: str) -> Optional[Dict[str, Any]]:
    """验证令牌"""
    try:
        return jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
    except PyJWTError as e:
        logger.error(f"令牌验证失败: {str(e)}")
        return None

def get_user_id_from_token(token: str) -> Optional[str]:
    """从令牌中获取用户ID"""
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        return payload.get("sub")
    except PyJWTError as e:
        logger.error(f"从令牌获取用户ID失败: {str(e)}")
        return None

def is_token_expired(token: str) -> bool:
    """检查令牌是否过期"""
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        exp = payload.get("exp")
        if exp:
            return datetime.utcnow() > datetime.fromtimestamp(exp)
        return True
    except PyJWTError:
        return True

def refresh_access_token(refresh_token: str) -> Optional[str]:
    """使用刷新令牌获取新的访问令牌"""
    try:
        payload = jwt.decode(refresh_token, SECRET_KEY, algorithms=[ALGORITHM])
        
        # 检查是否是刷新令牌
        if payload.get("type") != "refresh":
            logger.error("令牌类型错误，不是刷新令牌")
            return None
        
        user_id = payload.get("sub")
        if not user_id:
            logger.error("刷新令牌中缺少用户ID")
            return None
        
        # 创建新的访问令牌
        return create_access_token(user_id)
        
    except PyJWTError as e:
        logger.error(f"刷新访问令牌失败: {str(e)}")
        return None

# 为了兼容性，保留原有的jwt对象
class JWTManager:
    def __init__(self):
        self.secret_key = SECRET_KEY
        self.algorithm = ALGORITHM
    
    def encode(self, payload: Dict[str, Any]) -> str:
        return jwt.encode(payload, self.secret_key, algorithm=self.algorithm)
    
    def decode(self, token: str) -> Dict[str, Any]:
        return jwt.decode(token, self.secret_key, algorithms=[self.algorithm])

# 创建全局jwt实例
jwt_manager = JWTManager()

# 为了向后兼容，提供JWTError别名
JWTError = PyJWTError

# 兼容性别名
jwt = {
    'encode': encode,
    'decode': decode
}
